import os
import json

import numpy as np
import cv2 as cv
import torch
from torch.utils.data import Dataset

from utils.preprocess import median_blur

class Sen12SinglePatchDataset(Dataset):

    def __init__(self, config) -> None:
        super(Sen12SinglePatchDataset, self).__init__()

        self.dset_dir = config.base_dir
        with open(config.list_file, 'r') as fp:
            self.json_dict = json.load(fp)

        self.list = []
        for fn in self.json_dict.keys():
            self.list.append(
                os.path.join(self.dset_dir, fn)
            )


        print(f"Sen12 Single Patch Dataset created with {self.__len__()} pairs")

    def __getitem__(self, index) -> tuple:
        filename = self.list[index]
        patch = cv.imread(filename)
        fn = os.path.basename(filename)
        label, modal = self.json_dict[fn][0], self.json_dict[fn][1]

        if modal == "sar":
            patch = median_blur(patch)

        if len(patch.shape) <= 2:
            patch = np.expand_dims(patch, axis=2)

        if patch.shape[2] > 1:
            patch = patch.mean(axis=2, keepdims=True)

        patch = np.transpose(patch, [2, 0, 1])
        patch = torch.tensor(patch)

        return (patch, label)

    def __len__(self):
        return len(self.list)